package com.dequan.springaiproject.controller;

import com.dequan.springaiproject.manager.ChatClientManager;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;


/**
 *
 * 测试 LLM
 */
@RestController
@RequestMapping("/test")
public class TextChatLLM {

    @Autowired
    private ChatModel chatModel;

    @Autowired
    private ChatClient doctorChatClient;

    @Autowired
    private ChatClientManager chatClientManager;

    @Autowired
    private ChatClient withMemoryAdvisorChatClient;

    @Autowired
    private ChatClient withDBMemoryAdvisorChatClient;

    @GetMapping("/chat/model")
    public String chatModel(@RequestParam String message){
        if (message.isEmpty()) {
            return "请输入内容";
        }
        return chatModel.call(new Prompt(message)).getResult().getOutput().getText();
    }

    /**
     * 聊天流
     * @param message
     * @return
     */
    @GetMapping(value = "/chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chatStream(@RequestParam String message) {
        if (message.isEmpty()) {
            return Flux.just("请输入内容");
        }

        return doctorChatClient.prompt()
                .user(message)
                .stream()
                .content();
    }

    /**
     * SseEmitter 实现的聊天流
     * @param message
     * @return
     */
    @GetMapping(value = "/chat/streamWithEmitter", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatStreamWithEmitter(@RequestParam String message) {
        SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);

        try {
            Flux<String> stream = chatModel.stream(new Prompt(message))
                    .map(response -> response.getResult().getOutput().getText());

            stream.subscribe(
                    data -> {
                        try {
                            emitter.send(SseEmitter.event().data(data));
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    },
                    emitter::completeWithError,
                    emitter::complete
            );
        } catch (Exception e) {
            emitter.completeWithError(e);
        }

        return emitter;
    }

    /**
     * 轮训多模型测试
     * @param message
     * @return
     */
    @GetMapping("/chat/models")
    public String chatModels(@RequestParam String message){
        if (message.isEmpty()) {
            return "请输入内容";
        }

        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < 3; i++) {

            String template = "第 %d 次输出->\n";
            stringBuilder.append(String.format(template, i + 1));
            stringBuilder.append(
                    chatClientManager.getNextChatClient()
                            .prompt()
                            .user(message)
                            .call()
                            .content()
            );
            stringBuilder.append("\n");
        }
        return stringBuilder.toString();
    }


    @GetMapping("/chat/memory")
    public String chatMemory(@RequestParam String message) {
        if (message.isEmpty()) {
            return "请输入内容";
        }

        return withMemoryAdvisorChatClient.prompt(message).call().content();
    }

    @GetMapping("/chat/dbMemory")
    public String chatDbMemory(
            @RequestParam String message,
            @RequestParam String conversationId
    ) {
        if (message.isEmpty()) {
            return "请输入内容";
        }

        return withDBMemoryAdvisorChatClient.prompt()
                .user(message)
                .advisors(advisor -> advisor.param(ChatMemory.CONVERSATION_ID, conversationId))
                .call()
                .content();
    }

}